import java.nio.FloatBuffer;
import java.nio.IntBuffer;
import java.nio.LongBuffer;
import org.lwjgl.PointerBuffer;
import org.lwjgl.glfw.Callbacks;
import org.lwjgl.glfw.GLFWErrorCallback;
import org.lwjgl.glfw.GLFWKeyCallback;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.system.MemoryUtil;
import org.lwjgl.vulkan.*;
import static org.lwjgl.vulkan.KHRSurface.*;
import static org.lwjgl.vulkan.KHRSwapchain.*;
import static org.lwjgl.vulkan.VK11.*;
import static org.lwjgl.glfw.GLFW.*;
import static org.lwjgl.glfw.GLFWVulkan.*;

/**
 * Vulkan 1.1 的 Demo (未完成)
 * <li>glfw的基本使用;
 * <li>在桌面平台使用Vulkan的方法;
 * <li>最简单的Vertex Shader和Fragment Shader;
 * <li>与native库交换数据所用的高效的MemoryStack;
 * <li>完整的初始化, 循环, 销毁三大步骤;
 *
 * @see <a href="https://vulkan-tutorial.com/Drawing_a_triangle/Setup/Base_code">Drawing a triangle</a>
 * @see <a href="https://www.glfw.org/docs/latest/vulkan_guide.html">Vulkan guide</a>
 * @see <a href="https://github.com/LWJGL/lwjgl3/blob/master/modules/samples/src/test/java/org/lwjgl/demo/vulkan/HelloVulkan.java">HelloVulkan</a>
 */
public final class VulkanDemo {
	private static long window;
	private static VkInstance vkInst;
	private static VkDevice vkDev;
	private static long surface;
	private static VkQueue vkQueue;
	private static long cmdPool;
	private static VkCommandBuffer cmdBuf;
	private static long swapchain;
	private static SwapchainBuffers[] scBufs;
	private static long depth;
	private static long depthMem;
	private static long depthView;

	private static class SwapchainBuffers {
		long image;
		VkCommandBuffer cmd;
		long view;
	}

	private static void checkVkResult(int r) {
		if (r != VK_SUCCESS)
			throw new AssertionError("Vulkan error: " + r);
	}

	private static void init() {
		//noinspection resource
		GLFWErrorCallback.createPrint().set();
		if (!glfwInit())
			throw new AssertionError("glfwInit failed");
		if (!glfwVulkanSupported())
			throw new AssertionError("glfwVulkanSupported failed");
		glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
		window = glfwCreateWindow(512, 512, "LwjglVulkanDemo", 0, 0);
		if (window == 0)
			throw new AssertionError("glfwCreateWindow failed");
		//noinspection resource
		final GLFWKeyCallback cb = glfwSetKeyCallback(window, (win, key, _, action, _) -> {
			if (key == GLFW_KEY_ESCAPE && action == GLFW_RELEASE)
				glfwSetWindowShouldClose(win, true);
		});
		if (cb != null)
			cb.free();

		try (final MemoryStack stack = MemoryStack.stackPush()) {
			final IntBuffer outInt = stack.callocInt(1);
			final LongBuffer outLong = stack.callocLong(1);
			final PointerBuffer outPtr = stack.callocPointer(1);

			final VkApplicationInfo appInfo = VkApplicationInfo.calloc(stack)
					.sType$Default()
					.pApplicationName(stack.UTF8("DemoApp"))
					.applicationVersion(VK_MAKE_VERSION(1, 0, 0))
					.pEngineName(stack.UTF8("DemoEngine"))
					.engineVersion(VK_MAKE_VERSION(1, 0, 0))
					.apiVersion(VK_API_VERSION_1_1); // VK.getInstanceVersionSupported()
			final PointerBuffer exts = glfwGetRequiredInstanceExtensions();
			if (exts == null)
				throw new AssertionError("glfwGetRequiredInstanceExtensions failed");
			final VkInstanceCreateInfo icInfo = VkInstanceCreateInfo.calloc(stack)
					.sType$Default()
					.pApplicationInfo(appInfo)
					.ppEnabledExtensionNames(exts);
			checkVkResult(vkCreateInstance(icInfo, null, outPtr));
			vkInst = new VkInstance(outPtr.get(0), icInfo);

			checkVkResult(vkEnumeratePhysicalDevices(vkInst, outInt, null));
			final int devs = outInt.get(0);
			final PointerBuffer phyDevBuf = stack.callocPointer(devs);
			checkVkResult(vkEnumeratePhysicalDevices(vkInst, outInt, phyDevBuf));
			VkPhysicalDevice phyDev = null;
			int gfIdx = 0;
			for (int i = 0; i < devs; i++) {
				final VkPhysicalDevice dev = new VkPhysicalDevice(phyDevBuf.get(i), vkInst);
				final VkPhysicalDeviceProperties devProp = VkPhysicalDeviceProperties.calloc(stack);
				vkGetPhysicalDeviceProperties(dev, devProp);
				System.out.println("device: " + devProp.deviceNameString());
				if (phyDev == null) {
					vkGetPhysicalDeviceQueueFamilyProperties(dev, outInt, null);
					final int qfs = outInt.get(0);
					final VkQueueFamilyProperties.Buffer qfpBuf = VkQueueFamilyProperties.calloc(qfs, stack);
					vkGetPhysicalDeviceQueueFamilyProperties(dev, outInt, qfpBuf);
					for (int j = 0; j < qfs; j++) {
						if ((qfpBuf.get(j).queueFlags() & VK_QUEUE_GRAPHICS_BIT) != 0) {
							phyDev = dev;
							gfIdx = j;
							break;
						}
					}
				}
			}
			if (phyDev == null)
				throw new AssertionError("no graphics physical device");

			checkVkResult(glfwCreateWindowSurface(vkInst, window, null, outLong));
			surface = outLong.get(0);

			final VkDeviceQueueCreateInfo.Buffer qcInfo = VkDeviceQueueCreateInfo.calloc(1, stack)
					.sType$Default()
					.queueFamilyIndex(gfIdx)
					.pQueuePriorities(stack.floats(1.0f));
			final PointerBuffer extNames = stack.callocPointer(1);
			extNames.put(0, stack.ASCII(VK_KHR_SWAPCHAIN_EXTENSION_NAME));
			final VkPhysicalDeviceFeatures devFeat = VkPhysicalDeviceFeatures.calloc(stack);
			final VkDeviceCreateInfo dcInfo = VkDeviceCreateInfo.calloc(stack)
					.sType$Default()
					.pQueueCreateInfos(qcInfo)
					.ppEnabledExtensionNames(extNames)
					.pEnabledFeatures(devFeat);
			checkVkResult(vkCreateDevice(phyDev, dcInfo, null, outPtr));
			vkDev = new VkDevice(outPtr.get(0), phyDev, dcInfo);

			vkGetDeviceQueue(vkDev, gfIdx, 0, outPtr);
			vkQueue = new VkQueue(outPtr.get(0), vkDev);

			final VkCommandPoolCreateInfo cpcInfo = VkCommandPoolCreateInfo.calloc(stack)
					.sType$Default()
					.flags(VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT)
					.queueFamilyIndex(gfIdx);
			checkVkResult(vkCreateCommandPool(vkDev, cpcInfo, null, outLong));
			cmdPool = outLong.get(0);

			final VkCommandBufferAllocateInfo cmd = VkCommandBufferAllocateInfo.calloc(stack)
					.sType$Default()
					.commandPool(cmdPool)
					.level(VK_COMMAND_BUFFER_LEVEL_PRIMARY)
					.commandBufferCount(1);
			checkVkResult(vkAllocateCommandBuffers(vkDev, cmd, outPtr));
			cmdBuf = new VkCommandBuffer(outPtr.get(0), vkDev);

			final VkSwapchainCreateInfoKHR scInfo = VkSwapchainCreateInfoKHR.calloc(stack)
					.sType$Default()
					.surface(surface)
					.minImageCount(1)
					.imageFormat(VK_FORMAT_B8G8R8A8_UNORM)
					.imageColorSpace(KHRSurface.VK_COLOR_SPACE_SRGB_NONLINEAR_KHR)
					.imageExtent(VkExtent2D.calloc(stack)
							.width(512)
							.height(512))
					.imageArrayLayers(1)
					.imageUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT)
					.imageSharingMode(VK_SHARING_MODE_EXCLUSIVE)
					.preTransform(VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR)
					.compositeAlpha(VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR)
					.presentMode(VK_PRESENT_MODE_FIFO_KHR)
					.clipped(true);
			checkVkResult(vkCreateSwapchainKHR(vkDev, scInfo, null, outLong));
			swapchain = outLong.get(0);

			checkVkResult(vkGetSwapchainImagesKHR(vkDev, swapchain, outInt, null));
			final int scs = outInt.get(0);
			final LongBuffer scImgBuf = stack.callocLong(scs);
			checkVkResult(vkGetSwapchainImagesKHR(vkDev, swapchain, outInt, scImgBuf));

			scBufs = new SwapchainBuffers[scs];
			for (int i = 0; i < scs; i++) {
				scBufs[i] = new SwapchainBuffers();
				scBufs[i].image = scImgBuf.get(i);
				final VkImageViewCreateInfo ivcInfo = VkImageViewCreateInfo.calloc(stack)
						.sType$Default()
						.image(scBufs[i].image)
						.viewType(VK_IMAGE_VIEW_TYPE_2D)
						.format(VK_FORMAT_B8G8R8A8_UNORM)
						.components(it -> it
								.r(VK_COMPONENT_SWIZZLE_R)
								.g(VK_COMPONENT_SWIZZLE_G)
								.b(VK_COMPONENT_SWIZZLE_B)
								.a(VK_COMPONENT_SWIZZLE_A))
						.subresourceRange(it -> it
								.aspectMask(VK_IMAGE_ASPECT_COLOR_BIT)
								.baseMipLevel(0)
								.levelCount(1)
								.baseArrayLayer(0)
								.layerCount(1));
				checkVkResult(vkCreateImageView(vkDev, ivcInfo, null, outLong));
				scBufs[i].view = outLong.get(0);
			}

			final VkImageCreateInfo imcInfo = VkImageCreateInfo.calloc(stack)
					.sType$Default()
					.imageType(VK_IMAGE_TYPE_2D)
					.format(VK_FORMAT_D16_UNORM)
					.extent(it -> it
							.width(512)
							.height(512)
							.depth(1))
					.mipLevels(1)
					.arrayLayers(1)
					.samples(VK_SAMPLE_COUNT_1_BIT)
					.tiling(VK_IMAGE_TILING_OPTIMAL)
					.usage(VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT);

			/* create image */
			checkVkResult(vkCreateImage(vkDev, imcInfo, null, outLong));
			depth = outLong.get(0);

			final VkMemoryRequirements memReq = VkMemoryRequirements.calloc(stack);
			vkGetImageMemoryRequirements(vkDev, depth, memReq);
			final VkMemoryAllocateInfo maInfo = VkMemoryAllocateInfo.calloc(stack)
					.sType$Default()
					.allocationSize(memReq.size());
			checkVkResult(vkAllocateMemory(vkDev, maInfo, null, outLong));
			depthMem = outLong.get(0);

			checkVkResult(vkBindImageMemory(vkDev, depth, depthMem, 0));
			setImageLayout(depth, VK_IMAGE_ASPECT_DEPTH_BIT, VK_IMAGE_LAYOUT_UNDEFINED,
					VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL, 0);

			final VkImageViewCreateInfo ivcInfo = VkImageViewCreateInfo.calloc(stack)
					.sType$Default()
					.image(depth)
					.viewType(VK_IMAGE_VIEW_TYPE_2D)
					.format(VK_FORMAT_D16_UNORM)
					.subresourceRange(it -> it
							.aspectMask(VK_IMAGE_ASPECT_DEPTH_BIT)
							.baseMipLevel(0)
							.levelCount(1)
							.baseArrayLayer(0)
							.layerCount(1));

			checkVkResult(vkCreateImageView(vkDev, ivcInfo, null, outLong));
			depthView = outLong.get(0);

			var vertData = new float[]{
					// pos.xyz, color.rgb
					1, 1, 0, 1, 0, 0,
					-1, 1, 0, 0, 1, 0,
					0, -1, 0, 0, 0, 1
			};
			var idxData = new int[]{0, 1, 2};

			final VkBufferCreateInfo bcInfo = VkBufferCreateInfo.calloc(stack)
					.sType$Default()
					.size(vertData.length * 4)
					.usage(VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT)
					.sharingMode(VK_SHARING_MODE_EXCLUSIVE);

			checkVkResult(vkCreateBuffer(vkDev, bcInfo, null, outLong));
			final long vertBuf = outLong.get(0);

			vkGetBufferMemoryRequirements(vkDev, vertBuf, memReq);
			maInfo.allocationSize(memReq.size());
			checkVkResult(vkAllocateMemory(vkDev, maInfo, null, outLong));
			final long vertMem = outLong.get(0);

			checkVkResult(vkMapMemory(vkDev, vertMem, 0, maInfo.allocationSize(), 0, outPtr));
			final FloatBuffer data = outPtr.getFloatBuffer(0, (int)maInfo.allocationSize() >> 2);
			data.put(vertData).flip();
			vkUnmapMemory(vkDev, vertMem);
			checkVkResult(vkBindBufferMemory(vkDev, vertBuf, vertMem, 0));

			final VkVertexInputBindingDescription.Buffer vi_bindings = VkVertexInputBindingDescription.calloc(1);
			final VkVertexInputAttributeDescription.Buffer vi_attrs = VkVertexInputAttributeDescription.calloc(2);
			final VkPipelineVertexInputStateCreateInfo vi = VkPipelineVertexInputStateCreateInfo.calloc()
					.sType$Default()
					.pVertexBindingDescriptions(vi_bindings)
					.pVertexAttributeDescriptions(vi_attrs);

			vi_bindings.get(0)
					.binding(0)
					.stride(vertData.length / 3 * 4)
					.inputRate(VK_VERTEX_INPUT_RATE_VERTEX);
			vi_attrs.get(0)
					.binding(0)
					.location(0)
					.format(VK_FORMAT_R32G32B32_SFLOAT)
					.offset(0);
			vi_attrs.get(1)
					.binding(0)
					.location(1)
					.format(VK_FORMAT_R32G32_SFLOAT)
					.offset(4 * 3);

//			final VkDescriptorSetLayoutCreateInfo descriptor_layout = VkDescriptorSetLayoutCreateInfo.malloc(stack)
//					.sType$Default()
//					.pBindings(VkDescriptorSetLayoutBinding.calloc(1, stack)
//							.binding(0)
//							.descriptorType(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER)
//							.descriptorCount(DEMO_TEXTURE_COUNT)
//							.stageFlags(VK_SHADER_STAGE_FRAGMENT_BIT)
//					);
//
//			LongBuffer layouts = stack.mallocLong(1);
//			check(vkCreateDescriptorSetLayout(device, descriptor_layout, null, layouts));
//			desc_layout = layouts.get(0);
//
//			VkPipelineLayoutCreateInfo pPipelineLayoutCreateInfo = VkPipelineLayoutCreateInfo.calloc(stack)
//					.sType$Default()
//					.pSetLayouts(layouts);
//
//			check(vkCreatePipelineLayout(device, pPipelineLayoutCreateInfo, null, lp));
//			pipeline_layout = lp.get(0);

			VkAttachmentDescription.Buffer attachments = VkAttachmentDescription.calloc(2, stack);
			attachments.get(0)
					.flags(0)
					.format(VK_FORMAT_B8G8R8A8_UNORM)
					.samples(VK_SAMPLE_COUNT_1_BIT)
					.loadOp(VK_ATTACHMENT_LOAD_OP_CLEAR)
					.storeOp(VK_ATTACHMENT_STORE_OP_STORE)
					.stencilLoadOp(VK_ATTACHMENT_LOAD_OP_DONT_CARE)
					.stencilStoreOp(VK_ATTACHMENT_STORE_OP_DONT_CARE)
					.initialLayout(VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL)
					.finalLayout(VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);
			attachments.get(1)
					.flags(0)
					.format(VK_FORMAT_D16_UNORM)
					.samples(VK_SAMPLE_COUNT_1_BIT)
					.loadOp(VK_ATTACHMENT_LOAD_OP_CLEAR)
					.storeOp(VK_ATTACHMENT_STORE_OP_DONT_CARE)
					.stencilLoadOp(VK_ATTACHMENT_LOAD_OP_DONT_CARE)
					.stencilStoreOp(VK_ATTACHMENT_STORE_OP_DONT_CARE)
					.initialLayout(VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL)
					.finalLayout(VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL);

			VkSubpassDescription.Buffer subpass = VkSubpassDescription.calloc(1, stack)
					.pipelineBindPoint(VK_PIPELINE_BIND_POINT_GRAPHICS)
					.colorAttachmentCount(1)
					.pColorAttachments(VkAttachmentReference.calloc(1, stack)
							.attachment(0)
							.layout(VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL))
					.pDepthStencilAttachment(VkAttachmentReference.calloc(stack)
							.attachment(1)
							.layout(VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL));

			VkRenderPassCreateInfo rp_info = VkRenderPassCreateInfo.calloc(stack)
					.sType$Default()
					.pAttachments(attachments)
					.pSubpasses(subpass);

			checkVkResult(vkCreateRenderPass(vkDev, rp_info, null, outLong));
			long renderPass = outLong.get(0);

			final LongBuffer att = stack.longs(0, depthView);

			final VkFramebufferCreateInfo fcInfo = VkFramebufferCreateInfo.calloc(stack)
					.sType$Default()
					.renderPass(renderPass)
					.pAttachments(att)
					.width(512)
					.height(512)
					.layers(1);
			final LongBuffer frameBufs = MemoryUtil.memAllocLong(scBufs.length);
			for (int i = 0; i < scBufs.length; i++) {
				att.put(0, scBufs[i].view);
				checkVkResult(vkCreateFramebuffer(vkDev, fcInfo, null, outLong));
				frameBufs.put(i, outLong.get(0));
			}

			System.out.println("init OK!");
		}
	}

	private static VkCommandBuffer setupCmdBuf;

	private static void setImageLayout(long image, int aspectMask, int old_image_layout, int new_image_layout, int srcAccessMask) {
		try (MemoryStack stack = MemoryStack.stackPush()) {
			if (setupCmdBuf == null) {
				VkCommandBufferAllocateInfo cmd = VkCommandBufferAllocateInfo.calloc(stack)
						.sType$Default()
						.commandPool(cmdPool)
						.level(VK_COMMAND_BUFFER_LEVEL_PRIMARY)
						.commandBufferCount(1);

				final PointerBuffer outPtr = stack.callocPointer(1);
				checkVkResult(vkAllocateCommandBuffers(vkDev, cmd, outPtr));
				setupCmdBuf = new VkCommandBuffer(outPtr.get(0), vkDev);

				VkCommandBufferBeginInfo cmd_buf_info = VkCommandBufferBeginInfo.calloc(stack)
						.sType$Default()
						.pInheritanceInfo(null);
				checkVkResult(vkBeginCommandBuffer(setupCmdBuf, cmd_buf_info));
			}

			VkImageMemoryBarrier.Buffer image_memory_barrier = VkImageMemoryBarrier.calloc(1, stack)
					.sType$Default()
					.srcAccessMask(srcAccessMask)
					.oldLayout(old_image_layout)
					.newLayout(new_image_layout)
					.image(image)
					.subresourceRange(it -> it
							.aspectMask(aspectMask)
							.baseMipLevel(0)
							.levelCount(1)
							.baseArrayLayer(0)
							.layerCount(1));

			int src_stages = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
			int dest_stages = VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT;

			if (srcAccessMask == VK_ACCESS_HOST_WRITE_BIT) {
				src_stages = VK_PIPELINE_STAGE_HOST_BIT;
			}

			switch (new_image_layout) {
			case VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL:
				image_memory_barrier.dstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT);
				break;
			case VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL:
				image_memory_barrier.dstAccessMask(VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT);
				dest_stages = VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
				break;
			case VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL:
				/* Make sure any Copy or CPU writes to image are flushed */
				image_memory_barrier.dstAccessMask(VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_INPUT_ATTACHMENT_READ_BIT);
				dest_stages = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
				break;
			case VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL:
				image_memory_barrier.srcAccessMask(VK_ACCESS_MEMORY_READ_BIT);
				/* Make sure anything that was copying from this image has completed */
				image_memory_barrier.dstAccessMask(VK_ACCESS_TRANSFER_READ_BIT);
				dest_stages = VK_PIPELINE_STAGE_TRANSFER_BIT;
				break;
			case VK_IMAGE_LAYOUT_PRESENT_SRC_KHR:
				image_memory_barrier.dstAccessMask(VK_ACCESS_MEMORY_READ_BIT);
				break;
			}

			vkCmdPipelineBarrier(setupCmdBuf, src_stages, dest_stages, 0, null, null, image_memory_barrier);
		}
	}

	private static void exit() {
		//TODO
		if (vkDev != null) {
			vkDestroyDevice(vkDev, null);
			vkDev = null;
		}
		if (vkInst != null) {
			vkDestroyInstance(vkInst, null);
			vkInst = null;
		}
		if (window != 0) {
			Callbacks.glfwFreeCallbacks(window);
			glfwDestroyWindow(window);
			window = 0;
		}
		glfwTerminate();
		//noinspection resource
		final GLFWErrorCallback cb = glfwSetErrorCallback(null);
		if (cb != null)
			cb.free();
	}

	private static void loop() {
		while (!glfwWindowShouldClose(window)) {
			glfwPollEvents();
			//TODO
		}
	}

	public static void main(@SuppressWarnings("unused") String[] args) {
		try {
			init();
			loop();
		} finally {
			exit();
		}
	}
}
